# @Time : 2021-05-29 10:53 
# @Author : baldSun
# @File : Strategy.py 
# @Software: PyCharm

import DBTools
from logDB import LogDB
import tushare as ts
from datetime import datetime

ts.set_token('b2109fc6949181eb2f598eae02f4de44e028e7eb20ffc39503a7fd9e')
pro = ts.pro_api()


# 策略虚基类
class Strategy(object):
    def __init__(self, share_holding_limit, purchase_limit, share_holding_cycle, surplus_funding, name):
        # 持股上限， 单日购买上限，剩余金额
        self.share_holding_limit = share_holding_limit
        self.purchase_limit = purchase_limit
        self.surplus_funding = surplus_funding
        self.share_holding_cycle = share_holding_cycle
        self.name = name

    # 指定购买计划
    def make_purchase_plan(self, date):
        raise NotImplementedError('父类make_purchase_plan方法未实现')

    # 购买股票
    def buy_stock(self, date):
        # date格式转换：从2000-01-04转成20000104
        ts_date = date.replace('-', '')
        # 分配资金：平均分给这个策略购买的每支股票
        funding = self.surplus_funding//self.share_holding_limit
        # 选择可购买的股票
        i = 0
        plan = [] #最终敲定的计划
        for stock in self.purchase_plan:
            # 格式：pro.daily(ts_code='600000.SH', trade_date='20180701')
            df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
            # 该股票今日无数据，不可购买，删除这条记录
            if df.empty:
                print('今日不可购股：'+df['ts_code'][0])
                continue
            # 已经达到当日购买上限
            if i < self.purchase_limit:
                plan.append(stock)
            i = i + 1
        # 购买股票
        self.position = [] # 持仓队列
        a_position_stock = {} # 一条持仓记录
        for stock in plan:
            df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
            # 购股
            amount_of_hand = funding//(df['open'][0]*100) #可以买几手
            self.surplus_funding = self.surplus_funding - (amount_of_hand*df.open*100)
        #     写入日志
            stock_name = pro.stock_basic(ts_code=df['ts_code'][0])['name'][0]
            log = (self.name, date, df['ts_code'][0], stock_name, 'buy', amount_of_hand, df['open'][0])
            res = LogDB().insert_log(log)
            print(res)
        #     记入持仓队列
            a_position_stock['ts_code'] = df['ts_code'][0]
            a_position_stock['open'] = df['open'][0]
            a_position_stock['purchasing_date'] = date #买入日期
            a_position_stock['holdings'] = amount_of_hand #持仓量(手)
            a_position_stock['is_firmly_hold'] = False #是否坚定持有(达到止盈条件)
            self.position.append(a_position_stock)
        print('持仓情况：')
        print(self.position)

    # 计算持仓队列中所有股票的持有期收益率  Rhp = 持有期期末价值/持有期期初价值-1
    def cal_holding_yield(self, date):
        ts_date = date.replace('-', '')
        i = 0
        for stock in self.position:
            df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
            close_price = float(df['close'][0])
            yield1 = close_price / float(stock['open']) - 1
            self.position[i]['yield'] = yield1
            # print(self.position[i])
            i = i + 1

    # 执行售出操作
    def execute_sale(self, date, stock, sell_price):
        print('售出股票' + stock['ts_code'])
        # 更改手中现金
        self.surplus_funding = self.surplus_funding + stock['holdings']*100*sell_price
        # 写入日志
        stock_name = pro.stock_basic(ts_code=stock['ts_code'])['name'][0]
        log = (self.name, date, stock['ts_code'], stock_name, 'sell', stock['holdings'], sell_price)
        res = LogDB().insert_log(log)
        print(res)
        self.position.remove(stock)

    # 售出股票
    def sale_stock(self, date):
        raise NotImplementedError('父类sale_stock方法未实现')


# 子策略：ah溢价
class AhYiJia(Strategy):
    def __init__(self, strategy_id,surplus_funding):
        # 持股上限1支， 单日买入数1支， 持股周期5天
        Strategy.__init__(self, 1, 1, 5, surplus_funding, 'ah_yijia')
        self.strategy_id = strategy_id
        # 设置ah溢价策略的备用股票数为5
        self.spare_stock_num = 5

    # 指定购股计划
    def make_purchase_plan(self, date):
        db = DBTools.DBTools()
        self.purchase_plan = db.get_min_ah_stock(self.spare_stock_num, date)

    # # 购买股票
    # def buy_stock(self, date):
    #     # date格式转换：从2000-01-04转成20000104
    #     ts_date = date.replace('-', '')
    #     # 分配资金：平均分给这个策略购买的每支股票
    #     funding = self.surplus_funding//self.share_holding_limit
    #     # 选择可购买的股票
    #     i = 0
    #     plan = [] #最终敲定的计划
    #     for stock in self.purchase_plan:
    #         # 格式：pro.daily(ts_code='600000.SH', trade_date='20180701')
    #         df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
    #         # 该股票今日无数据，不可购买，删除这条记录
    #         if df.empty:
    #             print('今日不可购股：'+df['ts_code'][0])
    #             continue
    #         # 已经达到当日购买上限
    #         if i < self.purchase_limit:
    #             plan.append(stock)
    #         i = i + 1
    #     # 购买股票
    #     self.position = [] # 持仓队列
    #     a_position_stock = {} # 一条持仓记录
    #     for stock in plan:
    #         df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
    #         # 购股
    #         amount_of_hand = funding//(df['open'][0]*100) #可以买几手
    #         self.surplus_funding = self.surplus_funding - (amount_of_hand*df.open*100)
    #     #     写入日志
    #         stock_name = pro.stock_basic(ts_code=df['ts_code'][0])['name'][0]
    #         log = (self.name, date, df['ts_code'][0], stock_name, 'buy', amount_of_hand, df['open'][0])
    #         res = LogDB().insert_log(log)
    #         print(res)
    #     #     记入持仓队列
    #         a_position_stock['ts_code'] = df['ts_code'][0]
    #         a_position_stock['open'] = df['open'][0]
    #         a_position_stock['purchasing_date'] = date #买入日期
    #         a_position_stock['holdings'] = amount_of_hand #持仓量(手)
    #         a_position_stock['is_firmly_hold'] = False #是否坚定持有(达到止盈条件)
    #         self.position.append(a_position_stock)
    #     print('持仓情况：')
    #     print(self.position)
    # # 执行售出操作
    # def execute_sale(self, date, stock, sell_price):
    #     print('售出股票' + stock['ts_code'])
    #     # 更改手中现金
    #     self.surplus_funding = self.surplus_funding + stock['holdings']*100*sell_price
    #     # 写入日志
    #     stock_name = pro.stock_basic(ts_code=stock['ts_code'])['name'][0]
    #     log = (self.name, date, stock['ts_code'], stock_name, 'sell', stock['holdings'], sell_price)
    #     res = LogDB().insert_log(log)
    #     print(res)
    #     self.position.remove(stock)

    # # 计算持仓队列中所有股票的持有期收益率  Rhp = 持有期期末价值/持有期期初价值-1
    # def cal_holding_yield(self, date):
    #     ts_date = date.replace('-', '')
    #     i = 0
    #     for stock in self.position:
    #         df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
    #         close_price = float(df['close'][0])
    #         yield1 = close_price / float(stock['open']) - 1
    #         self.position[i]['yield'] = yield1
    #         # print(self.position[i])
    #         i = i + 1


    # 遍历持仓队列查找需要售出的股票
    def sale_stock(self, date):
        i = 0
        ts_date = date.replace('-', '')
        for stock in self.position:
            df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
            # 处理坚定持有的股票：
            if stock['is_firmly_hold']:
                if stock['yield'] <= 0.01:
                    self.execute_sale(date, stock, df['close'][0])
            # 处理常规股票
            else:
                # 达到持股周期, 卖出：
                purchasing_day = datetime.strptime(stock['purchasing_date'], '%Y-%m-%d')
                today = datetime.strptime(date, '%Y-%m-%d')
                if (today - purchasing_day).days >= self.share_holding_cycle:
                    self.execute_sale(date, stock, df['close'][0])
                # 达到止损条件，卖出：
                if stock['yield'] <= -0.3:
                    print(stock['ts_code'] + '达到止盈条件, ' + date + '抛售')
                    self.execute_sale(date, stock, df['close'][0])
                # 达到坚定持股条件，记录：
                if stock['yield'] >= 0.3:
                    print(stock['ts_code'] + '达到止盈条件, 从' + date + '开始坚定持股')
                    self.position[i]['is_firmly_hold'] = True
            i = i + 1


# 子策略：历史新高
class HighTop(Strategy):
    def __init__(self, strategy_id, surplus_funding):
        # 持股上限2支， 单日买入数1支， 持股周期5天
        Strategy.__init__(self, 2, 1, 10, surplus_funding, 'celue_hightop_high_high')
        self.strategy_id = strategy_id
        # 设置ah溢价策略的备用股票数为10
        self.spare_stock_num = 10

    # 指定购股计划
    def make_purchase_plan(self, date):
        db = DBTools.DBTools()
        self.purchase_plan = db.get_min_polocy_stock(self.spare_stock_num, date)

    # 遍历持仓队列查找需要售出的股票
    def sale_stock(self, date):
        i = 0
        ts_date = date.replace('-', '')
        for stock in self.position:
            df = pro.daily(ts_code=stock['ts_code'], trade_date=ts_date)
            # 处理坚定持有的股票：
            if stock['is_firmly_hold']:
                if stock['yield'] <= 0.05:
                    self.execute_sale(self, date, stock, df['close'][0])
            # 处理常规股票
            else:
                # 达到持股周期, 卖出：
                purchasing_day = datetime.strptime(stock['purchasing_date'], '%Y-%m-%d')
                today = datetime.strptime(date, '%Y-%m-%d')
                if (today - purchasing_day).days >= self.share_holding_cycle:
                    self.execute_sale(date, stock, df['close'][0])
                # 达到止损条件，卖出：
                if stock['yield'] <= -0.1:
                    print(stock['ts_code'] + '达到止损条件, 从' + date + '卖出')
                    self.execute_sale(date, stock, df['close'][0])
                # 达到坚定持股条件，记录：
                if stock['yield'] >= 0.8:
                    print(stock['ts_code'] + '达到止盈条件, 从' + date + '开始坚定持股')
                    self.position[i]['is_firmly_hold'] = True
            i = i + 1


if __name__ == '__main__':
    # 测试历史新高策略“购买股票”函数：
    high_top_obj = HighTop(0, 1000000)
    start_date = '1992-04-13'
    end_date = '1992-04-23'
    # 这几个函数必须严格按照下列顺序调用否则会报错
    # 因为前面的函数定义了后面需要使用的成员变量
    high_top_obj .make_purchase_plan(start_date)
    print()
    high_top_obj .buy_stock(start_date)
    print()
    high_top_obj .cal_holding_yield(end_date)
    print()
    # 测试卖出函数
    high_top_obj .sale_stock(end_date)
    print(high_top_obj .surplus_funding)
    # # 测试ah溢价策略“购买股票”函数：
    # ah_obj = AhYiJia(0, 1000000)
    # date = '2000-01-04'
    # # 这几个函数必须严格按照下列顺序调用否则会报错
    # # 因为前面的函数定义了后面需要使用的成员变量
    # ah_obj.make_purchase_plan('2000-01-04')
    # print()
    # ah_obj.buy_stock('2000-01-04')
    # print()
    # ah_obj.cal_holding_yield('2000-01-14')
    # print()
    # # 测试卖出函数
    # ah_obj.sale_stock('2000-01-14')
    # print(ah_obj.surplus_funding)